import mxnet as mx
import mxnet.ndarray as nd
from skimage import io
import numpy as np
from mxnet import recordio
import matplotlib.pyplot as plt
import os
from tqdm import tqdm

#将rec文件转img rec文件是MX框架的数据格式，InsightFace作者使用的是MX，所以采用了rec格式
#如果自己使用，有几种做法：其一是将rec转回到img 其二是直接使用rec，训练时将图片转为tensor格式即可
def rec2img(path_imgidx,path_imgrec,output_dir):
    # path_prefix = 'datasets/faces_webface_112x112/train'
    # path_imgidx = 'datasets/faces_webface_112x112/train.idx'
    # path_imgrec = 'datasets/faces_webface_112x112/train.rec'
    # output_dir = 'datasets/align_webface_112x112'

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    imgrec = recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r')


    for i in tqdm(range(501195)):
        header, s = recordio.unpack(imgrec.read_idx(i + 1))
        img = mx.image.imdecode(s).asnumpy()
        label = str(header.label)
        id = str(i)

        label_dir = os.path.join(output_dir, label)
        # 检查标签文件夹是否存在
        if not os.path.exists(label_dir):
            os.mkdir(label_dir)

        # plt.imshow(img)c
        # plt.title('id=' + str(i) + 'label=' + str(header.label))
        # plt.pause(0.1)
        # print('id=' + str(i) + 'label=' + str(header.label))

        fname = 'Figure_{}.png'.format(id)
        fpath = os.path.join(label_dir, fname)
        io.imsave(fpath, img)


def rec2img1(path_imgidx,path_imgrec,output_dir):
    # path_prefix = 'datasets/faces_webface_112x112/train'
    # path_imgidx = 'datasets/faces_webface_112x112/train.idx'
    # path_imgrec = 'datasets/faces_webface_112x112/train.rec'
    # output_dir = 'datasets/align_webface_112x112'

    imgrec = recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r')


    # header, s = recordio.unpack(imgrec.read_idx(490624 + 1))
    # print(s)
    # print(header.flag)
    # print(imgrec.read_idx(490624 + 1)[2*4:])
    # header, s = recordio.unpack(imgrec.read_idx(490000 + 1))
    # print(s)
    # print(header)
    for i in range(490000,len(imgrec.keys)):
        header, s = recordio.unpack(imgrec.read_idx(i + 1))
        print(i)
        if header.flag>0:
            s = imgrec.read_idx(i + 1)[4*header.flag:]
        img = mx.image.imdecode(s).asnumpy()


        label = str(header.label)
        id = str(i)

        if i %10000 == 0:
            print("read :",i)

        if img is  None:
            print("wrong image",label,id)

def read_rec_test(path_imgidx,path_imgrec):
    train_data = mx.image.ImageIter(batch_size=32,
                                    data_shape=(3, 112, 112),
                                    path_imgrec=path_imgrec,
                                    path_imgidx=path_imgidx,
                                    shuffle=False)

    print(train_data.imglist)
    train_data.reset()
    data_batch = train_data.next()
    data = data_batch.data[0]
    print(data_batch)
    print(data.shape)
    print(data[0].shape)
    plt.figure()
    for i in range(5):
        save_image = data[i].astype('uint8').asnumpy().transpose((1, 2, 0))
        plt.subplot(1, 5, i + 1)
        plt.imshow(save_image)
    plt.ion()


if __name__ == '__main__':

    rec2img1(r"e:\faces_webface_112x112\train.idx",r"e:\faces_webface_112x112\train.rec",r"e:\faces_webface_112x112\img")
    # read_rec_test(r"e:\faces_webface_112x112\train.idx",r"e:\faces_webface_112x112\train.rec")